import os
import numpy as np
import multiprocessing as mp


dst_path = '/home/lifuyu/data_disk/v-wewei_data/multi_obj_dataset_good_0326/image_tensors/tensors/'

def save_tensors(image_data, camera_intr_, camera_pose, k):
    #print('begin')
    np.save('/home/lifuyu/data_disk/v-wewei_data/depth/depth_{:05d}.npy'.format(k), image_data)
    np.save('/home/lifuyu/data_disk/v-wewei_data/camera_intr/camera_intr_{:05d}.npy'.format(k), camera_intr_)
    np.save('/home/lifuyu/data_disk/v-wewei_data/camera_pose/camera_pose_{:05d}.npy'.format(k), camera_pose)
    print('finished {}'.format(k))



if __name__ == "__main__":
    pool = mp.Pool(processes=mp.cpu_count())

    for i in range(0, 50):
        for j in range(2):
            depth_tensor = np.load(os.path.join(dst_path, 'depth_im_{:05d}.npz'.format(i * 2 + j)))['arr_0.npy']
            camera_intrs_tensor = np.load(os.path.join(dst_path, 'camera_intrs_{:05d}.npz'.format(i * 2 + j)))[
                'arr_0.npy']
            camera_pose_tensor = np.load(os.path.join(dst_path, 'camera_pose_{:05d}.npz'.format(i * 2 + j)))[
                'arr_0.npy']
            k = 200 * i + j * 100
            for image_data, camera_intr_, camera_pose in zip(depth_tensor, camera_intrs_tensor, camera_pose_tensor):
                pool.apply_async(save_tensors, args=(image_data, camera_intr_, camera_pose, k,))
                #save_tensors(image_data, camera_intr_, camera_pose, k)
                k += 1

    pool.close()
    pool.join()
